from torch.utils.data import Dataset
import pandas as pd
import numpy as np
from PIL import Image
from scipy import ndimage
import os
import cv2


class TrainDataset(Dataset):
    def __init__(self, csv_path='../Dataset/trainval.csv',
                 file_path='../Dataset/trainval/',
                 class_to_num=None, transform=None):
        self.file_path = file_path
        self.transform = transform
        self.class_to_num = class_to_num

        # read csv
        self.data_info = pd.read_csv(csv_path)  # ignore header
        # to ndarray
        self.image_arr = np.asarray(self.data_info.iloc[:, 0])  # shape (3000,)
        self.label_arr = np.asarray(self.data_info.iloc[:, 1])
        self.data_len = len(self.data_info.index)

    def __getitem__(self, index):
        # 从 image_arr中得到索引对应的文件名
        single_image = self.image_arr[index]
        # image = Image.open(os.path.join(self.file_path, single_image)).convert('RGB')

        image = cv2.imread(os.path.join(self.file_path, single_image))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        if self.transform is not None:
            image = self.transform(image=image)['image']  # albumentation transform

        label = self.label_arr[index]
        num_of_label = self.class_to_num[label]

        return image, num_of_label

    def __len__(self):
        return self.data_len


class TestDataset(Dataset):
    def __init__(self, csv_path='../Dataset/test.csv',
                 file_path='../Dataset/test/', transform=None):
        self.file_path = file_path
        self.transform = transform

        # read csv
        self.data_info = pd.read_csv(csv_path)
        # to ndarray
        self.image_arr = np.asarray(self.data_info.iloc[:, 0])  # image
        self.data_len = len(self.data_info.index)

    def __getitem__(self, index):
        # 从 image_arr中得到索引对应的文件名
        single_image = self.image_arr[index]

        # image = Image.open(self.file_path + single_image).convert('RGB')
        image = cv2.imread(os.path.join(self.file_path, single_image))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        if self.transform is not None:
            image = self.transform(image=image)['image']
            # image = self.transform(image)  # torchvison.transforms

        return image

    def __len__(self):
        return self.data_len


if __name__ == '__main__':
    data_info = pd.read_csv('../../Dataset-origin/trainval.csv')
    print(data_info.iloc[0:5, 1])
    print(len(data_info.index))
